/*=============================================================================
*
*   Copyright (C) 2020 All rights reserved.
*
*   Filename:		assemble_elemental_mat.hpp
*
*   Author: Wang Zhecheng - wangzhecheng@yeah.net
*
*   Date:  2020-02-27 18:29
*
*   Last Editors: Wang Zhecheng - wangzhecheng@yeah.net
*
*   Last modified:	2020-08-04 01:23
*
*   Description:
*
=============================================================================*/
#ifndef ASSEMBLE_BEE_H
#define ASSEMBLE_BEE_H

#include "../../geometry/shape/elemental_dof.hpp"

namespace geoxfem {
namespace model {

template <typename DimensionType, typename degree>
VectorXR
ShapeFunction(/*{{{*/
              const geometry::ShapeInterface<DimensionType, degree>& shp,
              const GaussPoint<DimensionType>&                       gp) {
  VectorXR N = shp.getStdNByLocal(gp.getLocalCoord());
  if(shp.ifXfem()) {
    VectorXR visible = gp.getVisibility();
    VectorXR M       = N.cwiseProduct(visible);
    VectorXR ans(N.size() * 2);
    ans << N, M;
    return ans;
  }
  return N;
} /*}}}*/

template <typename DimensionType, typename degree>
MatrixXR Deriv(/*{{{*/
               const geometry::ShapeInterface<DimensionType, degree>& shp,
               const GaussPoint<DimensionType>&                       gp) {
  size_t   nn        = shp.getNumberOfNeighbors(0);
  MatrixXR std_deriv = shp.getdNdgxByGlobal(gp.getCoord());
  if(shp.ifXfem()) {
    VectorXR visible    = gp.getVisibility();
    MatrixXR xfem_deriv = std_deriv;
    for(size_t i = 0; i != nn; ++i)
      xfem_deriv.col(i) *= visible(i);
    MatrixXR ans(std_deriv.rows(), std_deriv.cols() * 2);
    ans << std_deriv, xfem_deriv;
    return ans;
  }
  return std_deriv;
} /*}}}*/

struct assemble_bee { /*{{{*/
  template <typename DimensionType, typename degree>
  static MatrixXR
  apply(const geometry::ShapeInterface<DimensionType, degree>& shp,
        const GaussPoint<DimensionType>&                       gp) {
    return apply<traits::dimension<DimensionType>::value>(Deriv(shp, gp));
  }
  template <size_t dim>
  static MatrixXR apply(const MatrixXR&) {
    GEOXFEM_NOT_IMPLEMENTED();
    return MatrixXR{};
  }
  template <>
  MatrixXR apply<2>(const MatrixXR& dNdgx) {
    GEOXFEM_ASSERT(dNdgx.rows() == 2);
    auto     nn = dNdgx.cols();
    MatrixXR bee(3, nn * 2);

    for(auto i = 0; i != nn; ++i)
      bee.block(0, i * 2, 3, 2) << dNdgx(0, i), r0,  //
        r0, dNdgx(1, i),                             //
        dNdgx(1, i), dNdgx(0, i);
    /// check if has nan
    GEOXFEM_ASSERT(!bee.hasNaN());
    return bee;
  }
  template <>
  MatrixXR apply<3>(const MatrixXR& dNdgx) {
    GEOXFEM_ASSERT(dNdgx.rows() == 3);
    auto     nn = dNdgx.cols();
    MatrixXR bee(6, nn * 3);

    for(auto i = 0; i != nn; ++i)
      bee.block(0, i * 3, 6, 3) << dNdgx(0, i), r0, r0,  //
        r0, dNdgx(1, i), r0,                             //
        r0, r0, dNdgx(2, i),                             //
        dNdgx(1, i), dNdgx(0, i), r0,                    //
        dNdgx(2, i), r0, dNdgx(0, i),                    //
        r0, dNdgx(2, i), dNdgx(1, i);
    /// check if has nan
    GEOXFEM_ASSERT(!bee.hasNaN());
    return bee;
  }
};

template <typename DimensionType, typename degree>
MatrixXR BeeMatrix(const geometry::ShapeInterface<DimensionType, degree>& shp,
                   const GaussPoint<DimensionType>&                       gp) {
  return assemble_bee::apply(shp, gp);
}

/*}}}*/

//////////////////////////////////////////////////////////////
//////////////////////////////////////////////////////////////

template <typename GaussOrElem>
MatrixXR DeeMatrix(const GaussOrElem& obj) {
  const auto& stress = obj.getCurrSigma();
  const auto& param  = obj.getParam();
  const auto& mat    = traits::get_material<GaussOrElem>::apply(obj);
  MatrixXR    dee    = mat->getDeeMatrix(stress, param);
  GEOXFEM_ASSERT(!dee.hasNaN());
  return dee;
}

template <typename GaussOrElem>
MatrixXR ConductivityMatrix(const GaussOrElem& obj) {
  const auto mat = traits::get_material<GaussOrElem>::apply(obj);
  return mat->getSeepageTensor();
}

//////////////////////////////////////////////////////////////
//////////////////////////////////////////////////////////////

template <typename unknown1, typename unknown2>
struct assemble_gauss_kay { /*{{{*/
  template <typename DimensionType, typename degree>
  static MatrixXR
  apply(const geometry::ShapeInterface<DimensionType, degree>& shp,
        const GaussPoint<DimensionType>&                       gp) {
    return assemble_gauss_kay<unknown2, unknown1>::apply(shp, gp).transpose();
  }
  template <typename DimensionType, typename degree>
  MatrixXR
  operator()(const geometry::ShapeInterface<DimensionType, degree>& shp,
             const GaussPoint<DimensionType>&                       gp) const {
    return apply(shp, gp);
  }
};

/*}}}*/

template <typename unknown1, typename unknown2>
struct assemble_sub_kay { /*{{{*/
  template <typename DimensionType, typename degree>
  static MatrixXR
  apply(const geometry::ShapeInterface<DimensionType, degree>& shp) {
    return algorithm::integral(shp, assemble_gauss_kay<unknown1, unknown2>{});
  }
};
/*}}}*/

template <typename CouplePolicy>
struct assemble_elemental_kay { /*{{{*/
  template <typename DimensionType, typename degree, typename TimeStep>
  static MatrixXR apply(const geometry::ShapeInterface<DimensionType, degree>&,
                        const TimeStep&) {
    throw "not implemented yet";
  }
};
/*}}}*/

template <typename unknown1, typename unknown2>
struct assemble_sub_rhs_for_tstep { /*{{{*/
  template <typename DimensionType, typename degree>
  static VectorXR
  apply(const geometry::ShapeInterface<DimensionType, degree>& shp) {
    MatrixXR kay = assemble_sub_kay<unknown1, unknown2>::apply(shp);
    VectorXR ans = geometry::get_value_list_at_elemental_dof<
      unknown2, dof_value_at_prev_tstep>::apply(shp);
    return kay * ans;
  }
}; /*}}}*/

template <typename unknown1, typename unknown2>
struct assemble_sub_rhs_for_iter { /*{{{*/
  template <typename DimensionType, typename degree>
  static VectorXR
  apply(const geometry::ShapeInterface<DimensionType, degree>& shp) {
    MatrixXR kay   = assemble_sub_kay<unknown1, unknown2>::apply(shp);
    VectorXR d_ans = geometry::get_value_list_at_elemental_dof<
      unknown2, dof_value_at_curr_tstep>::apply(shp);
    return kay * d_ans;
  }
}; /*}}}*/

template <typename CouplePolicy>
struct assemble_element_rhs_for_tstep { /*{{{*/
  template <typename DimensionType, typename degree, typename TimeStep>
  static VectorXR apply(const geometry::ShapeInterface<DimensionType, degree>&,
                        const TimeStep&) {
    throw "not implemented yet";
  }
}; /*}}}*/

template <typename CouplePolicy>
struct assemble_element_rhs_for_iter { /*{{{*/
  template <typename DimensionType, typename degree, typename TimeStep>
  static VectorXR apply(const geometry::ShapeInterface<DimensionType, degree>&,
                        const TimeStep&) {
    throw "not implemented yet";
  }
}; /*}}}*/

template <typename CouplePolicy>
struct assemble_element { /*{{{*/
  template <typename DimensionType, typename degree, typename TimeStep>
  static void apply(const geometry::ShapeInterface<DimensionType, degree>& shp,
                    const TimeStep&             tstep,
                    Eigen::SparseMatrix<_real>& kay_mat,
                    VectorXR&                   rhs) {
    auto     dof_list = geometry::get_elemental_dof<CouplePolicy>::apply(shp);
    size_t   neq      = dof_list.size();
    VectorXR e_load;
    MatrixXR e_kay;
    e_load.setZero(neq);
    e_kay.setZero(neq, neq);

    /* get elemental kay matrix */
    e_kay = assemble_elemental_kay<CouplePolicy>::apply(shp, tstep);

    /* load by previous time step */
    e_load -= assemble_element_rhs_for_tstep<CouplePolicy>::apply(shp, tstep);

    /* load by current time step */
    e_load -= assemble_element_rhs_for_iter<CouplePolicy>::apply(shp, tstep);

    /* deal with dirichlet boundary */
    for(size_t i = 0; i != neq; ++i)
      if(dof_list.at(i)->ifDirichlet()) {
        /* deal kay matrix by dirichlet */
        VectorXR dirichlet_row(neq);
        dirichlet_row.setZero();
        dirichlet_row(i) = e_kay(i, i);
        e_kay.row(i)     = dirichlet_row.transpose();
        /* deal load by dirichlet */
        e_load[i] = dof_dirichlet_residual::apply(dof_list.at(i)) * e_kay(i, i);
      }

    Eigen::Index local_row = 0;
    for(auto& dof1 : dof_list) {
      Eigen::Index global_row = dof1->getIndex();
      /* set diagonal element at DOF with a dominant number at dirichlet */
      /* DOFs and set rhs at dirichlet DOFs at current iteration step */
      /* assamble global stiffness matrix at non-dirichlet DOFs */
      rhs[global_row] += e_load[local_row];
      Eigen::Index local_col = 0;
      for(auto& dof2 : dof_list) {
        Eigen::Index global_col = dof2->getIndex();
        kay_mat.coeffRef(global_row, global_col) += e_kay(local_row, local_col);
        ++local_col;
      }
      ++local_row;
    }
    return;
  }
}; /*}}}*/

}  // namespace model
}  // namespace geoxfem
#endif /* ASSEMBLE_H */
